Add Mooncake.jl extension#1343
Conversation
Port `MooncakeStaticArraysExt` from Mooncake.jl into StaticArrays.jl
as a package extension `StaticArraysMooncakeExt`. The extension teaches
Mooncake that `SArray{S,T,N,L}` with IEEE float / complex IEEE float
element types is a leaf-like differentiable container (its tangent
type is itself), and adds `frule!!` / `rrule!!` rules for `_new_` and
`getindex(::SArray, ::Int)`.
Tests live under `test/ext/mooncake/` with their own `Project.toml`
and a self-activating driver script, and run in a new `extra` CI job
(matrix-templated for future ext labels) on Julia 1.10 and current
stable.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
@ChrisRackauckas, this should fix SciML/SciMLSensitivity.jl#1425 |
|
Speaking from my experience maintaining Enzyme.jl, I would encourage this to be an extension in Mooncake for StaticArrays.jl StaticArrays has been very stable for a while, while the AD landscape is rapidly evolving. Looking back at the Enzyme extension in KernelAbstraction, it has caused the majority of patch version releases in the 0.9 series. Furthermore, the extension uses Mooncake internals quite a lot and barely requires special knowledge about StaticArrays, so it would make sense if this code lives where Mooncake experts can have an eye on it and there is less risk of it becoming stale. |
Why can it not be an extension in Mooncake? Extensions are perfectly symmetrical. |
|
Mooncake's policy means it covers Julia Base only. Existing pkg extensions maybe moved out in the future. Unfortunately I won't have time to read further comments. I'll fix review comments next week if any. Please feel free to close this if it is out of scope. |
That is your prerogative, but I would ask you to reconsider. A PR and an extension to another package is an implicit ask: "Please take on the maintenance of this piece of code." This is not a small ask, and the internals of an AD package are complicated. I have co-maintained Enzyme.jl, ForwardDiff.jl and even helped with Mooncake and it's interaction with the Julia compiler, yet when I look at this code, I am not confident that I could maintain it. When deciding where a package extension lives I try to look at the amount of internals of each package it is using. The current extension is mostly Mooncake and very little StaticArrays, so it would very naturally for me be a Mooncake extension. |
| @foldable function tangent_type(::Type{SArray{S,T,N,L}}) where {S,T<:_SElt,N,L} | ||
| return SArray{S,T,N,L} | ||
| end | ||
|
|
||
| @foldable function tangent_type( | ||
| ::Type{NoFData}, ::Type{SArray{S,T,N,L}} | ||
| ) where {S,T<:_SElt,N,L} | ||
| return SArray{S,T,N,L} | ||
| end | ||
|
|
||
| # Non-parametric alias used as a constraint, analogous to `CF` in complex.jl. | ||
| const _SAFloat = SArray{S,T,N,L} where {S,T<:_SElt,N,L} | ||
|
|
||
| @foldable fdata_type(::Type{T}) where {T<:_SAFloat} = NoFData | ||
| @foldable rdata_type(::Type{T}) where {T<:_SAFloat} = T | ||
|
|
||
| tangent(::NoFData, t::_SAFloat) = t |
There was a problem hiding this comment.
Aside from whether this should live in Mooncake or StaticArrays, was this even reviewed by a human before submission? These definitions look like they should be exactly identical to defaults.
Is this really the kind of pull requests you would welcome in your projects?
There was a problem hiding this comment.
Yes, this is actually needed -- the default tangent for Mooncake is not StaticArrays. Mooncake's automatically derived tangents are usually a NamedTuple of all field members of structs.
The tests for this PR run successfully locally, and I reviewed the code myself before creating it.
There was a problem hiding this comment.
If it's all necessary, then it's not much better. Nearly duplicating so much code for each array-like type is going to cause huge maintenance issues. Maybe it could all be generated by a macro that lives in Mooncake.jl?
Anyway, if you don't have the time to address issues raised in a review, making PRs is just a waste of time.
There was a problem hiding this comment.
If it's all necessary, then it's not much better. Nearly duplicating so much code for each array-like type is going to cause huge maintenance issues. Maybe it could all be generated by a macro that lives in Mooncake.jl?
Will Tebbutt primarily designed Moocanke’s tangent-type system while working with me. Automating this via metaprogramming is difficult, as subtypes of AbstractArray can vary considerably. This reflects lessons learned from Zygote, ChainRules, and other autograd systems. Will has nearly a decade of experience with automatic differentiation in Julia, so I trust his judgment on this.
if you don't have the time to address issues raised in a review, making PRs is just a waste of time.
The level of disagreement has made it difficult to sustain focused technical discussion.
There was a problem hiding this comment.
Automating this via metaprogramming is difficult, as subtypes of AbstractArray can vary considerably.
I have about 8 years of experience writing and maintaining Julia code, and in my opinion automation via metaprogramming would be the easier approach in the long run. This extension is its current form breaks major good engineering practices, for example not accessing internal functions of other packages and only using public API. Unless "internal" in function name is only decorative and they are not internal at all?
|
I'm closing this PR because I won't have time to continue discussions. The code here can serve as a template for anyone who wants to use StaticArrays with Mooncake. |
4 participants
StaticArraysMooncakeExtpackage extension porting Mooncake.jl's existingMooncakeStaticArraysExt. The extension teaches Mooncake thatSArray{S,T,N,L}with IEEE float / complex IEEE float element types is a leaf-like differentiable container (its tangent type is itself), and addsfrule!!/rrule!!rules forMooncake._new_(which IR normalisation lowersSArray(...)constructor calls to) andgetindex(::SArray, ::Int).test/ext/mooncake/with their ownProject.tomland a self-activating driver script, mirroring how Mooncake's own repo organises ext tests. A newextraCI job (matrix-templated for future ext labels) runs them on Julia 1.10 and the current stable.🤖 Generated with Claude Code